package com.youan.log.modules.audit.dataimport;

import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import com.youan.log.common.threadtask.ThreadTaskProgressor;
import com.youan.log.common.threadtask.ThreadTaskUtils;
import com.youan.log.modules.audit.dataanalysis.IP2Region;
import com.youan.log.modules.audit.dataimport.bean.TaskData;
import com.youan.log.modules.audit.entity.ImportRecord;
import com.youan.log.modules.audit.service.IImportRecordService;
import com.youan.log.modules.statistical.entity.LogStatistical;
import com.youan.log.modules.statistical.service.ILogStatisticalService;
import org.springblade.core.tool.utils.SpringUtil;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;

import static com.youan.log.modules.audit.dataimport.bean.ImportTaskKey.THREAD_TASK_STATUS_AGG;
import static com.youan.log.modules.audit.dataimport.bean.ImportTaskKey.THREAD_TASK_STATUS_AGG_TEXT;

/**
 * 聚合数据
 * @author xsh
 * @date 2021/3/8 14:39
 */
public class DataAggregator {

    private static final int COMMIT_TO_DB_SIZE = 10000;

    private TaskData taskData;
    private ThreadTaskProgressor taskProgress;
    private Long importRecordId;

    private IP2Region ip2Region ;
    private String countSql;
    private String aggSql;

    private JdbcTemplate jdbcTemplate;
    private List<LogStatistical> waitPushToDBList = new ArrayList<>(COMMIT_TO_DB_SIZE);
    private ILogStatisticalService logStatisticalService;

    public DataAggregator(TaskData taskData, ThreadTaskProgressor taskProgressor, Long importRecordId) {
        this.taskData = taskData;
        this.taskProgress = taskProgressor;
        this.importRecordId = importRecordId;

        DataSource dataSource = SpringUtil.getBean(DataSource.class);
        this.jdbcTemplate = new JdbcTemplate(dataSource);
        this.logStatisticalService = SpringUtil.getBean(ILogStatisticalService.class);

        this.ip2Region = new IP2Region();
        buildSql();
    }

    private void buildSql() {
        StringBuilder countSqlBuilder = new StringBuilder();
        countSqlBuilder.append("SELECT COUNT(1) FROM (SELECT COUNT(1) FROM ")
                .append(taskData.getTableName())
                .append(" GROUP BY rule_id, source_ip, destination_ip, date, direction) AS agg");
        this.countSql = countSqlBuilder.toString();

        StringBuilder aggSqlBuilder = new StringBuilder();
        aggSqlBuilder.append("SELECT date AS visit_date, direction AS visit_direction, source_ip, ")
                .append("destination_ip AS dest_ip, rule_id AS issue_type, sum(1) AS visit_number, min(CONCAT(date, ' ', time)) first_occurrence_time," +
                        " max(CONCAT(date, ' ', time)) latest_occurrence_time FROM ")
                .append(taskData.getTableName())
                .append(" GROUP BY rule_id, source_ip, destination_ip, date, direction");
        this.aggSql = aggSqlBuilder.toString();
    }

    public void aggregate() {
        updateTaskStatus();
        doAgg();
        this.ip2Region.close();
    }

    private void updateTaskStatus() {
        IImportRecordService importRecordService = SpringUtil.getBean(IImportRecordService.class);
        ImportRecord updateImportRecord = new ImportRecord();
        updateImportRecord.setId(this.importRecordId);
        updateImportRecord.setTaskStatus(THREAD_TASK_STATUS_AGG);
        importRecordService.updateById(updateImportRecord);

        int aggResultNum = getAggResultNum();
        this.taskProgress.resetProgressTotal(aggResultNum);

        ImportRecord importRecord = importRecordService.getById(this.importRecordId);
        ThreadTaskUtils.updateTaskStatus(importRecord.getTaskId(), THREAD_TASK_STATUS_AGG, THREAD_TASK_STATUS_AGG_TEXT);
    }

    private void doAgg() {
        LocalDate visitDate = LocalDate.parse(taskData.getDate(), DateTimeFormatter.ofPattern("yyyy-MM-dd"));

        this.jdbcTemplate.query(this.aggSql, rs -> {
            //SELECT date AS visit_date, direction AS visit_direction, source_ip, ")
            //                .append("destination_ip AS dest_ip, rule_id AS issue_type, sum(1) AS visit_number
            LogStatistical statistical = new LogStatistical();
            statistical.setId(IdWorker.getId());
            statistical.setVisitDate(visitDate);
            statistical.setVisitDirection(rs.getString("visit_direction"));
            statistical.setSourceIp(rs.getString("source_ip"));
            statistical.setDestIp(rs.getString("dest_ip"));
            statistical.setIssueType(rs.getString("issue_type"));
            statistical.setVisitNumber(rs.getInt("visit_number"));
            statistical.setFirstOccurrenceTime(rs.getString("first_occurrence_time"));
            statistical.setLatestOccurrenceTime(rs.getString("latest_occurrence_time"));
            // 生成源ip的国家和城市，目的ip的国家和城市
            IP2Region.Region region = ip2Region.ip2Region(rs.getString("source_ip"));
            if(region != null) {
                statistical.setSourceCountry(region.getCountry());
                statistical.setSourceProvince(region.getProvince());
                statistical.setSourceCity(region.getCity());
            }
            region = ip2Region.ip2Region(rs.getString("dest_ip"));
            if(region != null) {
                statistical.setDestCountry(region.getCountry());
                statistical.setDestProvince(region.getProvince());
                statistical.setDestCity(region.getCity());
            }
            waitPushToDBList.add(statistical);
            if(waitPushToDBList.size() >= COMMIT_TO_DB_SIZE) {
                pushToDB();
            }
        });
        pushToDB();
    }

    private void pushToDB() {
        this.logStatisticalService.saveBatch(this.waitPushToDBList);
        this.taskProgress.addProgressNum(this.waitPushToDBList.size());
        this.waitPushToDBList = new ArrayList<>(COMMIT_TO_DB_SIZE);
    }


    private int getAggResultNum() {
        int count = this.jdbcTemplate.queryForObject(countSql, Integer.class);
        return count;
    }


}
